import streamlit as st
import pandas as pd
import json
import os
from sqlalchemy import create_engine, text

from sqlalchemy import inspect
from sdv.single_table import GaussianCopulaSynthesizer
from sdv.single_table import CTGANSynthesizer
from sdv.metadata import Metadata


last_default_gen_num = 0

# 初始化数据源存储（模拟数据库）
DATA_SOURCES_FILE = "data_sources.json"
if not os.path.exists(DATA_SOURCES_FILE):
    with open(DATA_SOURCES_FILE, "w") as f:
        json.dump([], f)


# 加载数据源列表
def load_data_sources():
    with open(DATA_SOURCES_FILE, "r") as f:
        return json.load(f)


# 保存数据源列表
def save_data_sources(sources):
    with open(DATA_SOURCES_FILE, "w") as f:
        json.dump(sources, f)


# 页面布局
st.set_page_config(page_title="数据源管理系统", layout="wide")

# 初始化会话状态
if "current_page" not in st.session_state:
    st.session_state.current_page = "数据源管理"
if "edit_source" not in st.session_state:
    st.session_state.edit_source = None

# 侧边栏菜单
with st.sidebar:
    st.header("导航菜单")
    menu_option = st.radio(
        "选择功能",
        ["数据源管理", "快速导入",'多表关联导入'],
        index=0 if st.session_state.current_page == "数据源管理" else 1
    )
    st.session_state.current_page = menu_option

# ------------------------- 数据源管理页面 -------------------------
# ...（前面的代码保持不变，以下是需要修改的部分）

# ------------------------- 数据源管理页面 -------------------------
if st.session_state.current_page == "数据源管理":
    st.title("数据源管理")
    data_sources = load_data_sources()

    # 显示现有数据源
    st.subheader("已配置的数据源")
    if not data_sources:
        st.info("暂无数据源，请点击下方按钮添加")
    else:
        # 表头
        cols = st.columns([2, 2, 2, 1, 1])
        cols[0].markdown("**名称**")
        cols[1].markdown("**类型**")
        cols[2].markdown("**连接信息**")
        cols[3].markdown("**操作**")
        cols[4].markdown("**操作**")

        for idx, source in enumerate(data_sources):
            cols = st.columns([2, 2, 2, 1, 1])

            # 名称可点击（按钮形式）
            if cols[0].button(
                    f"📁 {source['name']}",
                    key=f"name_{idx}",
                    use_container_width=True
            ):
                st.session_state.selected_source = source
                st.session_state.show_tables = True  # 触发显示表列表

            cols[1].write(source["type"])
            cols[2].code(f"Host: {source.get('host', 'N/A')}\nDB: {source.get('database', 'N/A')}")

            # 编辑按钮
            if cols[3].button("✏️ 编辑", key=f"edit_{idx}"):
                st.session_state.edit_source = source

            # 删除按钮
            if cols[4].button("🗑️ 删除", key=f"del_{idx}"):
                data_sources.pop(idx)
                save_data_sources(data_sources)
                st.rerun()

    # ------------------------- 添加/编辑数据源表单 -------------------------
    with st.form("data_source_form"):
        st.subheader("添加新数据源" if not st.session_state.edit_source else "编辑数据源")

        # 表单字段
        source_name = st.text_input("数据源名称",
                                    value=st.session_state.edit_source.get("name",
                                                                           "") if st.session_state.edit_source else "")
        source_type = st.selectbox(
            "类型",
            ["MySQL", "PostgreSQL", "CSV", "Excel"],
            index=0 if not st.session_state.edit_source else ["MySQL", "PostgreSQL", "CSV", "Excel"].index(
                st.session_state.edit_source["type"])
        )

        # 数据库配置字段
        if source_type in ["MySQL", "PostgreSQL"]:
            col1, col2 = st.columns(2)
            host = col1.text_input("主机地址",
                                   value=st.session_state.edit_source.get("host",
                                                                          "") if st.session_state.edit_source else "")
            port = col2.number_input("端口",
                                     value=st.session_state.edit_source.get("port",
                                                                            3306) if st.session_state.edit_source else 3306)
            database = st.text_input("数据库名",
                                     value=st.session_state.edit_source.get("database",
                                                                            "") if st.session_state.edit_source else "")
            username = st.text_input("用户名",
                                     value=st.session_state.edit_source.get("username",
                                                                            "") if st.session_state.edit_source else "")
            password = st.text_input("密码", type="password",
                                     value=st.session_state.edit_source.get("password",
                                                                            "") if st.session_state.edit_source else "")

        # 文件配置字段
        if source_type in ["CSV", "Excel"]:
            file_path = st.text_input("文件路径",
                                      value=st.session_state.edit_source.get("file_path",
                                                                             "") if st.session_state.edit_source else "")

        # 表单提交按钮
        submit_cols = st.columns([4, 1])
        if submit_cols[0].form_submit_button("💾 保存配置"):
            new_source = {
                "name": source_name,
                "type": source_type
            }
            if source_type in ["MySQL", "PostgreSQL"]:
                new_source.update({
                    "host": host,
                    "port": port,
                    "database": database,
                    "username": username,
                    "password": password
                })
            else:
                new_source["file_path"] = file_path

            # 更新或新增数据源
            if st.session_state.edit_source:
                # 更新现有数据源
                index = next(i for i, s in enumerate(data_sources) if s["name"] == st.session_state.edit_source["name"])
                data_sources[index] = new_source
            else:
                # 添加新数据源
                data_sources.append(new_source)

            save_data_sources(data_sources)
            st.session_state.edit_source = None
            st.rerun()

        # 取消编辑按钮
        if st.session_state.edit_source and submit_cols[1].form_submit_button("❌ 取消"):
            st.session_state.edit_source = None
            st.rerun()

    # ------------------------- 显示表列表及预览 -------------------------
    # ...（原有表列表预览代码保持不变）...
elif st.session_state.current_page == '快速导入':
    st.title("快速导入数据")
    data_sources = load_data_sources()

    # ------------------------- 源数据选择 -------------------------
    st.subheader("1. 选择数据源")
    col_source1, col_source2 = st.columns(2)
    source_type = col_source1.radio("源类型", ["数据库"], horizontal=True)
    selected_source = col_source2.selectbox("选择源数据源", [s["name"] for s in data_sources])

    if selected_source:
        source = next(s for s in data_sources if s["name"] == selected_source)
        try:
            if source_type == "数据库" and source["type"] in ["MySQL", "PostgreSQL"]:
                # ...（连接数据库和获取表名部分保持不变）...

                # ------------------------- 新增：全选/反选按钮 -------------------------
                st.subheader("2. 批量配置参数")
                col_batch1, col_batch2, col_batch3 = st.columns([2, 2, 4])
                default_generate_rows = col_batch1.number_input("默认生成行数", min_value=1, value=1000)  # ✅ 默认100条


                from sqlalchemy import inspect

                engine = create_engine(f"{'mysql+pymysql' if source['type'] == 'MySQL' else 'postgresql'}://"
                                       f"{source['username']}:{source['password']}@{source['host']}:"
                                       f"{source['port']}/{source['database']}")
                # 获取所有表名
                inspector = inspect(engine)
                all_tables = inspector.get_table_names()
                # 全选/反选按钮
                with col_batch2:
                    if st.button("全选"):
                        for table in all_tables:
                            st.session_state[f"check_{table}"] = True
                    if st.button("反选"):
                        for table in all_tables:
                            current = st.session_state.get(f"check_{table}", False)
                            st.session_state[f"check_{table}"] = not current

                # ------------------------- 表格配置区域 -------------------------
                st.subheader("3. 表格配置（勾选需要处理的表）")

                # 表格标题
                cols = st.columns([1, 2, 3, 3, 2])
                cols[0].markdown("**勾选**")
                cols[1].markdown("**原表名称**")
                cols[2].markdown("**训练条数**")
                cols[3].markdown("**目标表名称**")
                cols[4].markdown("**生成数量**")

                # 初始化 session_state
                if "table_configs" not in st.session_state:
                    st.session_state.table_configs = {}

                # 每行配置
                for table in all_tables:
                    cols = st.columns([1, 2, 3, 3, 2])

                    # 勾选框（自动初始化）
                    checked = cols[0].checkbox(
                        "",
                        value=st.session_state.table_configs.get(table, {}).get("checked", False),
                        key=f"check_{table}"
                    )

                    # 原表名称
                    cols[1].write(table)

                    # 训练条数（自动统计）
                    train_count = pd.read_sql(f"SELECT COUNT(*) FROM {table}", engine).iloc[0, 0]
                    cols[2].code(f"{train_count} 条")

                    # 目标表名称（带默认值）
                    target_table = cols[3].text_input(
                        "",
                        value=table,  # 默认与源表同名
                        key=f"target_{table}",
                        label_visibility="collapsed"
                    )

                    # 生成数量（使用默认值100）
                    gen_rows = cols[4].number_input(
                        "",
                        min_value=1,
                        value=st.session_state.table_configs.get(table, {}).get("gen_rows", default_generate_rows),
                        # ✅ 默认100
                        key=f"rows_{table}",
                        label_visibility="collapsed"
                    )

                    # 保存配置
                    st.session_state.table_configs[table] = {
                        "checked": checked,
                        "target_table": target_table,
                        "gen_rows": gen_rows
                    }

                # ------------------------- 目标数据配置 -------------------------
                st.subheader("4. 目标数据设置")
                col_target1, col_target2 = st.columns(2)
                target_source = col_target1.selectbox("选择目标数据源", [s["name"] for s in data_sources])
                write_mode = col_target2.radio("写入模式", ["覆盖", "追加"], horizontal=True)

                if st.button("🚀 开始生成", use_container_width=True):
                    target_config = next(s for s in data_sources if s["name"] == target_source)
                    progress_bar = st.progress(0)
                    selected_tables = [t for t, cfg in st.session_state.table_configs.items() if cfg["checked"]]
                    total = len(selected_tables)

                    for idx, table in enumerate(selected_tables):
                        try:
                            cfg = st.session_state.table_configs[table]

                            # 读取源数据
                            real_data = pd.read_sql_table(table, engine)
                            metadata = Metadata.detect_from_dataframe(real_data)
                            print(metadata)
                            # 训练模型（示例使用高斯Copula）
                            model = GaussianCopulaSynthesizer(metadata)
                            model.fit(real_data)

                            # 生成数据
                            synthetic_data = model.sample(cfg["gen_rows"])

                            # 连接目标数据库
                            target_engine = create_engine(
                                f"mysql+pymysql://{target_config['username']}:{target_config['password']}@"
                                f"{target_config['host']}:{target_config['port']}/{target_config['database']}"
                                if target_config["type"] == "MySQL"
                                else f"postgresql://{target_config['username']}:{target_config['password']}@"
                                     f"{target_config['host']}:{target_config['port']}/{target_config['database']}"
                            )

                            if write_mode == "覆盖":
                                # 检查表是否存在
                                if inspect(target_engine).has_table(cfg["target_table"]):
                                    # 清空数据但保留表结构
                                    with target_engine.connect() as conn:
                                        conn.execute(text(f"TRUNCATE TABLE {cfg['target_table']}"))  # 更高效的清空方式
                                        conn.commit()

                                # 使用append模式写入（保证不重建表）
                                synthetic_data.to_sql(
                                    name=cfg["target_table"],
                                    con=target_engine,
                                    if_exists="append",
                                    index=False

                                )
                            else:
                                # 直接追加数据
                                synthetic_data.to_sql(
                                    name=cfg["target_table"],
                                    con=target_engine,
                                    if_exists="append",
                                    index=False
                                )

                            # 更新进度
                            progress = (idx + 1) / total
                            progress_bar.progress(progress)
                            st.success(f"表 [{table}] → [{cfg['target_table']}] 生成成功！")
                        except Exception as e:
                            st.error(f"表 [{table}] 处理失败: {str(e)}")

                    st.balloons()
        except Exception as e:
            st.error(f"连接失败: {str(e)}")


elif st.session_state.current_page == "多表关联导入":
    st.title("多表关联导入")
    data_sources = load_data_sources()

    # 步骤1：选择源数据源
    selected_source = st.selectbox(
        "选择源数据源",
        [s["name"] for s in data_sources if s["type"] in ["MySQL", "PostgreSQL"]]
    )

    if selected_source:
        source = next(s for s in data_sources if s["name"] == selected_source)
        try:
            # 连接数据库
            engine = create_engine(
                f"mysql+pymysql://{source['username']}:{source['password']}@"
                f"{source['host']}:{source['port']}/{source['database']}"
                if source["type"] == "MySQL" else
                f"postgresql://{source['username']}:{source['password']}@"
                f"{source['host']}:{source['port']}/{source['database']}"
            )
            inspector = inspect(engine)

            # 步骤2：选择关联表
            st.subheader("选择关联表")
            all_tables = inspector.get_table_names()
            selected_tables = st.multiselect(
                "勾选需要关联的表（按主表→从表顺序选择）",
                all_tables,
                format_func=lambda x: f"📊 {x}"
            )

            if len(selected_tables) >= 2:
                # 步骤3：动态配置外键关系
                st.subheader("配置外键关联")
                st.caption("💡 请为每个从表选择其主表和关联字段")
                relationships = []

                # 获取所有表的字段信息
                table_columns = {}
                for table in selected_tables:
                    cols = inspector.get_columns(table)
                    table_columns[table] = [col["name"] for col in cols]
                    # 校验主键存在
                    pk = inspector.get_pk_constraint(table)["constrained_columns"]
                    if not pk:
                        st.error(f"表 {table} 未定义主键，无法建立关联！")
                        raise ValueError(f"表 {table} 缺少主键")

                # 为每个从表配置关系
                for i in range(1, len(selected_tables)):  # 从第二个表开始配置
                    child_table = selected_tables[i]

                    # 列布局
                    col1, col2, col3 = st.columns([2, 2, 3])
                    with col1:
                        # 只能选择前面的表作为父表
                        parent_table = st.selectbox(
                            f"为 {child_table} 选择主表",
                            selected_tables[:i],  # 关键修复：只允许选择前面的表
                            key=f"parent_{child_table}"
                        )

                    with col2:
                        # 显示主表主键
                        parent_pk = inspector.get_pk_constraint(parent_table)["constrained_columns"][0]
                        st.markdown(f"主表主键：\n`{parent_pk}`")

                    with col3:
                        # 动态获取子表字段
                        fk_options = table_columns[child_table]
                        default_fk = f"{parent_table.split('_')[0]}_id"  # 智能默认值
                        selected_fk = st.selectbox(
                            f"选择 {child_table} 的外键字段",
                            options=fk_options,
                            index=fk_options.index(default_fk) if default_fk in fk_options else 0,
                            key=f"fk_{child_table}"
                        )

                    relationships.append({
                        "parent": parent_table,
                        "child": child_table,
                        "foreign_key": selected_fk,
                        "parent_pk": parent_pk
                    })

                # 步骤4：生成配置
                st.subheader("生成配置")
                col1, col2 = st.columns(2)
                with col1:
                    base_rows = st.number_input("基础生成行数", min_value=100, value=1000)
                with col2:
                    target_source = st.selectbox("目标数据源", [s["name"] for s in data_sources])

                # 步骤5：生成数据
                if st.button("🚀 生成关联数据", type="primary"):
                    from sdv.relational import HMA1

                    # 构建元数据
                    metadata = {
                        "tables": {table: {"primary_key": inspector.get_pk_constraint(table)["constrained_columns"][0]}
                                   for table in selected_tables},
                        "relationships": [
                            {
                                "parent_table": rel["parent"],
                                "child_table": rel["child"],
                                "foreign_key": rel["foreign_key"]
                            } for rel in relationships
                        ]
                    }

                    # 训练模型
                    model = HMA1(metadata)
                    model.fit(engine)

                    # 生成数据
                    synthetic_data = model.sample(num_rows=base_rows)

                    # 写入目标数据库
                    target_config = next(s for s in data_sources if s["name"] == target_source)
                    target_engine = create_engine(
                        f"mysql+pymysql://{target_config['username']}:{target_config['password']}@"
                        f"{target_config['host']}:{target_config['port']}/{target_config['database']}"
                        if target_config["type"] == "MySQL" else
                        f"postgresql://{target_config['username']}:{target_config['password']}@"
                        f"{target_config['host']}:{target_config['port']}/{target_config['database']}"
                    )

                    # 按顺序写入数据
                    success_count = 0
                    for table in selected_tables:
                        try:
                            data = synthetic_data[table]
                            data.to_sql(
                                name=table,
                                con=target_engine,
                                if_exists="replace",
                                index=False
                            )
                            success_count += 1
                            st.success(f"表 {table} 写入成功（{len(data)} 行）")
                        except Exception as e:
                            st.error(f"表 {table} 写入失败: {str(e)}")

                    if success_count == len(selected_tables):
                        st.balloons()
        except Exception as e:
            st.error(f"操作失败: {str(e)}")